package com.ruoyi;

import com.ruoyi.common.constant.MilvusConstants;
import com.ruoyi.common.core.domain.entity.MilvusSchema;
import com.ruoyi.common.utils.MilvusUtils;
import com.ruoyi.common.utils.StringUtils;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.data.neo4j.Neo4jDataAutoConfiguration;
import org.springframework.boot.autoconfigure.data.neo4j.Neo4jRepositoriesAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.cache.annotation.EnableCaching;
import org.springframework.core.env.ConfigurableEnvironment;
import org.springframework.core.env.Environment;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.List;

/**
 * 启动程序
 *
 * @author ruoyi
 */
@SpringBootApplication(exclude = {
        DataSourceAutoConfiguration.class,
        Neo4jDataAutoConfiguration.class,
        Neo4jRepositoriesAutoConfiguration.class
})
@EnableCaching
@Slf4j
public class RuoYiApplication {


    public static void main(String[] args) {
        // System.setProperty("spring.devtools.restart.enabled", "false");
        ConfigurableEnvironment environment = SpringApplication.run(RuoYiApplication.class, args).getEnvironment();
        try {
            createMilvusCollection();
        } catch (Exception e) {
            log.error("milvus 初始化失败", e);
        }
        logApplicationStartup(environment);

    }

    /**
     * 启动程序时候判断垂直集合是否存在如果不存在就创建垂直集合
     */
    private static void createMilvusCollection() throws InterruptedException {

        //判断集合是否存在
        if (!MilvusUtils.hasCollection("Vector_Document_Segment")) {
            List<MilvusSchema> milvusSchemas = new ArrayList<>();

            //主键 (采用雪花算法生成id)
            MilvusSchema id = new MilvusSchema();
            id.setFieldName("id");
            id.setDescription("id");
            id.setAutoID(false);
            id.setPartitionKey(false);
            id.setPrimaryKey(true);
            id.setDataType(DataType.Int64.getCode());
            milvusSchemas.add(id);

            //生成一个稠密向量
            MilvusSchema vector = new MilvusSchema();
            vector.setFieldName("vector");
            vector.setDescription("稠密向量");
            vector.setAutoID(false);
            vector.setPartitionKey(false);
            vector.setPrimaryKey(false);
            vector.setDataType(DataType.FloatVector.getCode());
            vector.setDimension(1024);
            milvusSchemas.add(vector);
            //bm25 的稀疏向量
            MilvusSchema sparse = new MilvusSchema();
            sparse.setFieldName("sparse_vector");
            sparse.setDescription("稀疏向量");
            sparse.setAutoID(false);
            sparse.setPartitionKey(false);
            sparse.setPrimaryKey(false);
            sparse.setDataType(DataType.SparseFloatVector.getCode());
            milvusSchemas.add(sparse);
            //文本值
            MilvusSchema text = new MilvusSchema();
            text.setFieldName("text");
            text.setDescription("原文");
            text.setAutoID(false);
            text.setPartitionKey(false);
            text.setPrimaryKey(false);
            text.setMaxLength(10000);
            text.setDataType(DataType.VarChar.getCode());
            milvusSchemas.add(text);

            //JSON (用于存储文件名 文件元数据等)
            MilvusSchema json = new MilvusSchema();
            json.setFieldName("metadata");
            json.setDescription("向量元数据");
            json.setAutoID(false);
            json.setPartitionKey(false);
            json.setPrimaryKey(false);
            json.setMaxLength(10000);
            json.setDataType(DataType.JSON.getCode());
            milvusSchemas.add(json);

            CreateCollectionReq.CollectionSchema schema = MilvusUtils.createSchema(milvusSchemas);
            List<IndexParam> index = MilvusUtils.createIndex();
            MilvusUtils.createCollection(MilvusConstants.VECTOR_DOCUMENT_SEGMENT, schema, index);
            //睡眠两秒后加载集合
            Thread.sleep(2000);
            MilvusUtils.loadCollection(MilvusConstants.VECTOR_DOCUMENT_SEGMENT);

        }

    }

    private static void logApplicationStartup(Environment env) {
        String protocol = "http";
        if (env.getProperty("server.ssl.key-store") != null) {
            protocol = "https";
        }
        String serverPort = env.getProperty("server.port");
        String contextPath = env.getProperty("server.servlet.context-path");
        if (StringUtils.isBlank(contextPath)) {
            contextPath = "doc.html";
        } else {
            contextPath = contextPath + "doc.html";
        }
        String hostAddress = "localhost";
        try {
            hostAddress = InetAddress.getLocalHost().getHostAddress();
        } catch (UnknownHostException e) {
            log.warn("The host name could not be determined, using `localhost` as fallback");
        }
        log.info("""
                        \n----------------------------------------------------------
                        \t应用程序“{}”正在运行中......
                        \t接口文档访问 URL:
                        \t本地: \t\t{}://localhost:{}{}
                        \t外部: \t{}://{}:{}{}
                        \t配置文件: \t{}
                        ----------------------------------------------------------""",
                env.getProperty("llm-info.name"),
                protocol,
                serverPort,
                contextPath,
                protocol,
                hostAddress,
                serverPort,
                contextPath,
                env.getActiveProfiles());
    }
}
